import numpy as np
import pandas as pd
import cv2
import json
from tqdm import tqdm
import os
import sys
from PIL import Image
from skimage import morphology

mulMaskRoot = "/home/eziao/datasets/AliYun/results"
binMaskRoot = "/home/eziao/datasets/AliYun/road/results"
newMaskRoot = '/home/eziao/datasets/AliYun/newResults'


# TODO 骨架提取
def boneExtract(imagePath):

    img = cv2.imread(imagePath, 0)
    _, binary = cv2.threshold(img, 200, 255, cv2.THRESH_BINARY_INV)
    binary[binary == 255] = 1
    skel, distance = morphology.medial_axis(binary, return_distance=True)
    dist_on_skel = distance * skel
    dist_on_skel = dist_on_skel.astype(np.uint8)*255


# 判断是否为 1,2 像素值的道路标注
def maskCheck(image):
    if np.sum(image == 1) < 3:  # 核心
        return False
    return True


def maskRussianVote(maskMul, maskBin):  # LOL
    newMask = np.where(maskBin != 0, 4, maskMul)  # 道路
    return newMask


# def erode(image):
if __name__ == '__main__':
    for image in tqdm(os.listdir(mulMaskRoot)):

        mulMaskPath = os.path.join(mulMaskRoot, image)
        mulMask = np.array(Image.open(mulMaskPath))

        binMaskPath = os.path.join(binMaskRoot, image)
        (binMaskName, ext) = os.path.splitext(binMaskPath)
        if ext == ".png":
            binMask = np.array(Image.open(binMaskPath))
            if maskCheck(binMask):
                predictMask = maskRussianVote(mulMask, binMask)
            else:
                predictMask = mulMask

        newMaskName = os.path.splitext(image)[0] + '.png'

        predictMask = Image.fromarray(predictMask)
        predictMask.save(os.path.join(newMaskRoot, newMaskName))
